import torch
from .ranger import Ranger
import logging


def make_optimizer(cfg, model, center_criterion):
    logger = logging.getLogger("reid_baseline.train")
    params = []
    for key, value in model.named_parameters():
        if not value.requires_grad:
            continue
        lr = cfg.SOLVER.BASE_LR
        weight_decay = cfg.SOLVER.WEIGHT_DECAY
        if "bias" in key:
            lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR
            weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS
        if cfg.SOLVER.LARGE_FC_LR:
            if "classifier" in key or "arcface" in key:
                "对FC层使用双倍学习率？"
                lr = cfg.SOLVER.BASE_LR * 2
                logger.info('Using two times learning rate for fc ')
        if "gap" in key:
            lr = cfg.SOLVER.BASE_LR * 10
            weight_decay = 0

        params += [{"params": [value], 'initial_lr': lr, "lr": lr, "weight_decay": weight_decay}]
    if cfg.SOLVER.OPTIMIZER_NAME == 'SGD':
        optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params, momentum=cfg.SOLVER.MOMENTUM)
    elif cfg.SOLVER.OPTIMIZER_NAME == 'Ranger':
        optimizer = Ranger(params)
        logger.info('using Ranger for optimizer ')
    else:
        optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params)
    "对center loss单独使用SGD优化器"
    if isinstance(center_criterion, list):
        center_criterion_params = []
        for _center_criterion in center_criterion:
            center_criterion_params.append({'params': _center_criterion.parameters()})
            optimizer_center = torch.optim.SGD(params=center_criterion_params, lr=cfg.SOLVER.CENTER_LR)
    else:
        optimizer_center = torch.optim.SGD(params=center_criterion.parameters(), lr=cfg.SOLVER.CENTER_LR)

    return optimizer, optimizer_center
